#define DEBUG_TYPE "ABC"

#include <string>
#include <sstream>

#include "llvm/BasicBlock.h"
#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/GlobalVariable.h"
#include "llvm/InstrTypes.h"
#include "llvm/Instructions.h"
#include "llvm/LLVMContext.h"
#include "llvm/Module.h"
#include "llvm/Type.h"
#include "llvm/Value.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/GetElementPtrTypeIterator.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "LLVMCompat.h"

using namespace llvm;

STATISTIC(ABCCounter, "Total number of array bound checks added");

namespace {

struct ABCPass : public FunctionPass {

  static char ID;
  static const StringRef out_bound_msg;
  Function* printf_func;
  Function* exit_func;
  Value* message;
  unsigned version;

  ABCPass() : FunctionPass(ID) { }

  virtual bool doInitialization(Module &M){
    //init statisitcs
    ABCCounter = 0;

    //init version
    version = 0;

    //define a global variable for out_bound_msg
    Constant* string_const =
    		#if (LLVM_MAJOR > 3) || ((LLVM_MAJOR == 3) && (LLVM_MINOR > 0))
	    	         ConstantDataArray::getString(M.getContext(),out_bound_msg);
	    	 #else
	        ConstantArray::get(M.getContext(), out_bound_msg);
	        #endif
    message = new GlobalVariable(M, string_const->getType(), true,
                                 GlobalVariable::InternalLinkage, string_const,
                                 "OutOfBoundMsg");
    //create printf function
    //  signature: i32 @printf(i8*, ...)
    std::vector<Type*> params;
    params.push_back(Type::getInt8PtrTy(M.getContext()));
    FunctionType* printf_type = FunctionType::get(Type::getInt32Ty(M.getContext()),
                                                  params, true);
    printf_func = cast<Function>(
        M.getOrInsertFunction("printf", printf_type));

    //create exit function
    params.clear();
    params.push_back(Type::getInt32Ty(M.getContext()));

    FunctionType* exit_type = FunctionType::get(Type::getVoidTy(M.getContext()), 
                                                params, false);
    exit_func = cast<Function>(M.getOrInsertFunction("exit", exit_type));
    return true;
  }

  inline std::string IntToString(int i) {
    std::stringstream ss;
    std::string ret;
    ss << i;
    ret = ss.str();
    return ret;
  }

  //create a new basic block where check fails
  BasicBlock* createCheckFailBlock(Function* F, Value* V, Value* Sz) {
    //create printf call instruction
    std::vector<Value*> idxlist;
    Constant* const_zero = ConstantInt::get(Type::getInt32Ty(F->getContext()), 0);
    idxlist.push_back(const_zero);
    idxlist.push_back(const_zero);
    std::vector<Value*> args;
    args.push_back(ConstantExpr::getGetElementPtr(cast<Constant>(message), idxlist, true));
    Constant* const_checknum = ConstantInt::get(Type::getInt32Ty(F->getContext()), version);
    args.push_back(const_checknum);
    args.push_back(V);
    args.push_back(Sz);
    CallInst* printf_call = CallInst::Create(printf_func, args, "call_printf");
    //create exit call instruction
    Constant* const_one = ConstantInt::get(Type::getInt32Ty(F->getContext()), 1);
    args.clear();
    args.push_back(const_one);
    CallInst* exit_call = CallInst::Create(exit_func, args);
    exit_call->addAttribute(~0, Attribute::NoReturn);
    exit_call->addAttribute(~0, Attribute::NoUnwind);
    //add unreachable instruction after exit
    UnreachableInst* end_instr = new UnreachableInst(F->getContext()); 
    //create return instruciton
   // Constant* const_one = ConstantInt::get(Type::getInt32Ty(F->getContext()), 1);
   // ReturnInst* ret_instr = ReturnInst::Create(F->getContext(), const_one);

    //create basic block
    BasicBlock* BB = BasicBlock::Create(F->getContext(), 
                                        "ABC_Check_Fail" + IntToString(version), F);
    //BB->getInstList().insert(BB->begin(), ret_instr);
    BB->getInstList().insert(BB->begin(), end_instr);
    BB->getInstList().insert(BB->begin(), exit_call);
    BB->getInstList().insert(BB->begin(), printf_call);

    version ++;

    return BB;
  }

  virtual bool runOnFunction(Function &F) {
    bool changed = false;

    std::vector<GetElementPtrInst*> gep_instrs;
    for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I) {
      if(GetElementPtrInst* gep_instr = dyn_cast<GetElementPtrInst>(&(*I))) {
        gep_instrs.push_back(gep_instr);
      } /* end if GEP instruction */
    } /* end for each instructions */

    if (gep_instrs.size() == 0) return changed;

    for (std::vector<GetElementPtrInst*>::iterator I = gep_instrs.begin(),
         E = gep_instrs.end(); I != E; I ++) {
      GetElementPtrInst* gep_instr = (*I);
      //iter through the gepi types
      gep_type_iterator GI = gep_type_begin(gep_instr), GE = gep_type_end(gep_instr);
      //position in GI
      int pos = 0;
      while (GI != GE) {
        //if array type
        if (const ArrayType* array = dyn_cast<ArrayType>(*GI)) {
          changed = true;
          //get bound and the value
          unsigned size = array->getNumElements();
          Value* value = gep_instr->getOperand(pos + 1);
          if ((size == 0) || (value == NULL)) {
            GI ++;
            continue;
          }
          Value* bound = dyn_cast<Value>(
              ConstantInt::get(value->getType(), size));

          //split the block before GEP instruction
          BasicBlock* curr_block = gep_instr->getParent();
          BasicBlock* SB = curr_block->splitBasicBlock(gep_instr);
          BasicBlock* FB = createCheckFailBlock(&F, value, bound);
          //insert compare instruction before the terminator
          ICmpInst* cmp_instr = new ICmpInst(
              curr_block->getTerminator(), ICmpInst::ICMP_SLT, value, bound);
          //insert a branch instruciton at the end of current block
          BranchInst* br_instr = BranchInst::Create(SB, FB, cmp_instr);
          //replace the original terminator with the branch
          llvm::ReplaceInstWithInst(curr_block->getTerminator(), br_instr);

          ABCCounter ++;
        } /* if array type */
        GI ++;
        pos ++;
      } /* end while gep type */
    } /* end for each gep */

    return changed;
  }

  virtual bool doFinalization(Module &M) {
    errs() << "[ABC]>>>Total number of ABC checks added: "
        << ABCCounter << '\n';
    return false;
  }
};
const StringRef ABCPass::out_bound_msg = "Array Bound Check Failed:%d, V:%d, Sz:%d\n";
char ABCPass::ID = 0; RegisterPass<ABCPass> X("abc", "Array Bounds Checking");
}
